from Agent.ReAct import ReActAgent
from Models.Factory import ChatModelFactory
from Tools import *
from Tools.RagTool import RagTool
from langchain_community.chat_message_histories.in_memory import ChatMessageHistory

import glob, os


def launch_agent(agent: ReActAgent, welcome_message: str = "有什么可以帮您？"):
    human_icon = "\U0001f468"
    ai_icon = "\U0001f916"
    chat_history = ChatMessageHistory()
    first_turn = True
    while True:
        if not first_turn:
            task = input(f"{human_icon}：")
        else:
            task = input(f"{ai_icon}：{welcome_message}\n{human_icon}：")
            first_turn = False

        if task.strip() == "":
            continue

        if task.strip().lower() == "quit":
            break
        reply = agent.run(task, chat_history, verbose=True)
        print(f"{ai_icon}：{reply}\n")


def main():
    # llm
    llm = ChatModelFactory.get_model("qwen-max")

    # 定义工具集
    document_path = "ragdata/docs/"
    index_path = "ragdata/index/"
    build_index_required = False
    files = glob.glob(os.path.join(document_path, "*.jsonl"))
    tools = [
        RagTool(
            rag_id=file.replace(".jsonl", ""),
            document_path=file,
            index_path=os.path.join(index_path, file.replace(".jsonl", "")),
            build_index_required=build_index_required,
        ).as_rag_tool()
        for file in files
    ]
    tools.append(get_datetime_tool)
    tools.append(finish_placeholder)
    # tools = [
    #    RagTool(
    #        rag_id="湖南长远锂科股份有限公司",
    #        document_path = os.path.join(document_path, "湖南长远锂科股份有限公司.jsonl"),
    #        index_path = os.path.join(index_path, "湖南长远锂科股份有限公司"),
    #        build_index_required = build_index_required
    #    ).as_tool(),
    #    finish_placeholder
    # ]

    # 定义智能体
    agent = ReActAgent(
        llm=llm,
        tools=tools,
        work_dir="./data",
        main_prompt_file="./prompts/AgenticRag.txt",
        max_thought_steps=10,
    )

    # 运行智能体
    launch_agent(agent)


if __name__ == "__main__":
    main()
